import random
import tensorflow as tf


# 小批量读取样本
def data_iter(batch_size, sample, labels):
    """
    batch_size:批量大小
    sample:样本数据集
    labels:结果集（标签集）
    这段代码定义了一个名为 data_iter 的函数，它接受三个参数：batch_size，sample 和 labels。函数的目的是将数据集分成多个批次，每个批次的大小为 batch_size。

首先，它获取样本数据集的长度 sample_len，并根据样本长度定义一个下标数组 indices。

然后，它使用 random.shuffle 函数随机打乱下标数组。

接下来，它使用一个 for 循环遍历下标数组，每次取出一个大小为 batch_size 的批次。在循环体内部，它使用 tf.constant 和 tf.gather 函数从样本数据集和标签集中获取对应的批次数据，并使用 yield 语句返回这些数据。

这段代码的目的是将数据集分成多个随机批次，可以用于训练神经网络模型。
    """
    sample_len = len(sample)  # 获取到样本数组长度
    indices = list(range(sample_len))  # 根据样本长度定义数组下标

    # 随机读取小批量样本
    random.shuffle(indices) # 打乱数组下标
    for i in range(0, sample_len, batch_size):
        j = tf.constant(indices[i:min(i+batch_size, sample_len)])
        yield tf.gather(sample, j), tf.gather(labels, j)


# 模型定义
def linreg(X, w, b):
    """
    这段代码定义了一个名为 linreg 的函数，它接受三个参数：X，w 和 b。函数的目的是计算线性回归模型的输出。

具体来说，它使用矩阵乘法计算 X 和 w 的乘积，并将结果加上偏置项 b。最后，它返回计算结果。

这段代码实现了线性回归模型的预测公式：y = Xw + b。
    """
    return tf.matmul(X, w) + b


# 定义损失函数
def squared_loss(y_hat, y):
    """均方损失 (y_hat-y)**EXP2
    这段代码定义了一个名为squared_loss的函数，它用来计算均方损失。函数接受两个参数：y_hat和y。其中，y_hat表示预测值，而y表示真实值。

在函数体中，我们首先使用tf.reshape函数将y调整为与y_hat相同的形状。然后，我们计算预测值与真实值之差的平方，并除以2。这样，我们就得到了均方损失。

均方损失是一种常用的损失函数，它用来衡量预测值与真实值之间的差异。在回归问题中，它经常被用来作为优化目标。

    """
    return (y_hat - tf.reshape(y, y_hat.shape))**2 / 2


# 定义优化算法
def sgd(params, grads, lr, batch_size):
    """小批量梯度下降
    这段代码是一个小批量梯度下降算法的实现。它接受四个参数：params，grads，lr和batch_size。其中，params是一个包含模型参数的列表，grads是一个包含梯度的列表，lr是学习率，而batch_size是批量大小。

在函数体中，我们使用了一个for循环来遍历模型参数和梯度。对于每一对参数和梯度，我们使用assign_sub方法来更新参数。这个方法会将参数减去学习率乘以梯度除以批量大小的结果。

这样，我们就可以使用这个函数来更新模型参数，从而实现小批量梯度下降算法。
    """
    for params, grads in zip(params, grads):
        params.assign_sub(lr*grads/batch_size)


